package priv.yiyuan.example.aesencryption;

import javax.crypto.*;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.io.*;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.util.Base64;

/**
 * 流式文件加密工具类
 * 使用AES/CBC/PKCS5Padding算法进行流式加密和解密
 * 返回处理后的输入流
 */

/**
 * @author YiYuan
 * @data 2025/8/13
 * @apoNote
 */
public class StreamEncryptionUtil {

    private static final String ALGORITHM = "AES";

    private static final String TRANSFORMATION = "AES/CBC/PKCS5Padding";

    private static final int KEY_SIZE = 256;

    private static final int IV_SIZE = 16;

    private static final int BUFFER_SIZE = 4096;

    /**
     * 生成AES密钥
     *
     * @return Base64编码的密钥
     */
    public static String generateKey() throws NoSuchAlgorithmException {
        KeyGenerator keyGenerator = KeyGenerator.getInstance(ALGORITHM);
        keyGenerator.init(KEY_SIZE);
        SecretKey secretKey = keyGenerator.generateKey();
        return Base64.getEncoder().encodeToString(secretKey.getEncoded());
    }

    /**
     * 加密输入流并返回加密后的输入流
     *
     * @param inputStream 待加密的输入流
     * @param key         Base64编码的AES密钥
     * @return 加密后的输入流
     */
    public InputStream encryptStream(InputStream inputStream, String key)
            throws NoSuchAlgorithmException, NoSuchPaddingException,
            InvalidKeyException, InvalidAlgorithmParameterException {

        // 将Base64密钥转换为SecretKey
        byte[] decodedKey = Base64.getDecoder().decode(key);
        SecretKey secretKey = new SecretKeySpec(decodedKey, 0, decodedKey.length, ALGORITHM);

        // 生成随机IV
        byte[] iv = new byte[IV_SIZE];
        SecureRandom random = new SecureRandom();
        random.nextBytes(iv);
        IvParameterSpec ivParameterSpec = new IvParameterSpec(iv);

        // 初始化加密器
        Cipher cipher = Cipher.getInstance(TRANSFORMATION);
        cipher.init(Cipher.ENCRYPT_MODE, secretKey, ivParameterSpec);

        // 创建管道流
        PipedInputStream pipedIn = new PipedInputStream();
        try {
            PipedOutputStream pipedOut = new PipedOutputStream(pipedIn);

            // 在新线程中处理加密
            new Thread(() -> {
                try (OutputStream out = pipedOut) {
                    // 首先写入IV
                    out.write(iv);

                    // 使用CipherOutputStream进行加密
                    try (CipherOutputStream cipherOut = new CipherOutputStream(out, cipher)) {
                        byte[] buffer = new byte[BUFFER_SIZE];
                        int bytesRead;
                        while ((bytesRead = inputStream.read(buffer)) != -1) {
                            cipherOut.write(buffer, 0, bytesRead);
                        }
                    }
                } catch (IOException e) {
                    throw new RuntimeException("Encryption failed", e);
                } finally {
                    try {
                        inputStream.close();
                    } catch (IOException e) {
                        // Ignore
                    }
                }
            }).start();

            return pipedIn;
        } catch (IOException e) {
            throw new RuntimeException("Failed to create pipe", e);
        }
    }

    /**
     * 解密输入流并返回解密后的输入流
     *
     * @param encryptedStream 加密的输入流
     * @param key             Base64编码的AES密钥
     * @return 解密后的输入流
     */
    public InputStream decryptStream(InputStream encryptedStream, String key)
            throws NoSuchAlgorithmException, NoSuchPaddingException,
            InvalidKeyException, InvalidAlgorithmParameterException, IOException {

        // 将Base64密钥转换为SecretKey
        byte[] decodedKey = Base64.getDecoder().decode(key);
        SecretKey secretKey = new SecretKeySpec(decodedKey, 0, decodedKey.length, ALGORITHM);

        // 读取IV
        byte[] iv = new byte[IV_SIZE];
        int bytesRead = encryptedStream.read(iv);
        if (bytesRead != IV_SIZE) {
            throw new IOException("Invalid encrypted stream format");
        }
        IvParameterSpec ivParameterSpec = new IvParameterSpec(iv);

        // 初始化解密器
        Cipher cipher = Cipher.getInstance(TRANSFORMATION);
        cipher.init(Cipher.DECRYPT_MODE, secretKey, ivParameterSpec);

        // 创建管道流
        PipedInputStream pipedIn = new PipedInputStream();
        try {
            PipedOutputStream pipedOut = new PipedOutputStream(pipedIn);

            // 在新线程中处理解密
            new Thread(() -> {
                try (OutputStream out = pipedOut;
                     CipherInputStream cipherIn = new CipherInputStream(encryptedStream, cipher)) {
                    byte[] buffer = new byte[BUFFER_SIZE];
                    int bytesReadThread;
                    while ((bytesReadThread = cipherIn.read(buffer)) != -1) {
                        out.write(buffer, 0, bytesReadThread);
                    }
                } catch (IOException e) {
                    throw new RuntimeException("Decryption failed", e);
                }
            }).start();

            return pipedIn;
        } catch (IOException e) {
            throw new RuntimeException("Failed to create pipe", e);
        }
    }

    /**
     * 加密字节数组并返回输入流
     *
     * @param data 待加密的数据
     * @param key  Base64编码的AES密钥
     * @return 加密后的输入流
     */
    public InputStream encryptBytes(byte[] data, String key)
            throws NoSuchAlgorithmException, NoSuchPaddingException,
            InvalidKeyException, InvalidAlgorithmParameterException {
        return encryptStream(new ByteArrayInputStream(data), key);
    }

    /**
     * 解密字节数组并返回输入流
     *
     * @param encryptedData 加密的数据
     * @param key           Base64编码的AES密钥
     * @return 解密后的输入流
     */
    public InputStream decryptBytes(byte[] encryptedData, String key)
            throws NoSuchAlgorithmException, NoSuchPaddingException,
            InvalidKeyException, InvalidAlgorithmParameterException, IOException {
        return decryptStream(new ByteArrayInputStream(encryptedData), key);
    }

    public static void main(String[] args) throws NoSuchAlgorithmException {
        System.out.println(StreamEncryptionUtil.generateKey());
    }
}
